package errorgroup

import (
	"github.com/mars/pkg/loggers"
	"github.com/mars/pkg/trace"
)

// Group 允许用户提交一个或多个任务，然后等待任务返回的结果.
// 使用时候，请务必保证任务的数量和指定的size一致
// 每个Group实例只可以执行一次**Complete**方法
type Group struct {
	size    int
	results chan interface{}
}

func WithSize(size int) *Group {
	return &Group{
		results: make(chan interface{}, size),
		size:    size,
	}
}

func UnSize() *Group {
	return &Group{
		results: make(chan interface{}, 1),
	}
}

func (g *Group) Go(f func() interface{}) *Group {
	go func() {
		defer func() {
			if err := recover(); err != nil {
				loggers.Errorf("Go execute panic: %v, stack: %s", err, trace.Stack(2))
				close(g.results)
			}
		}()
		g.results <- f()
	}()
	return g
}

func (g *Group) GoBatch(fns ...func() interface{}) *Group {
	for _, f := range fns {
		g.Go(f)
	}
	return g
}

func (g *Group) WhenComplete(completeFn func(result interface{})) {
	completeFn(<-g.results)
}

func (g *Group) Wait() interface{} {
	return <-g.results
}

func (g *Group) WhenCompleteAsync(completeFn func(result interface{})) {
	go func() {
		defer func() {
			if err := recover(); err != nil {
				loggers.Errorf("WhenCompleteAsync execute panic: %+v, stack: %s", err, string(trace.Stack(2)))
			}
		}()
		g.WhenComplete(completeFn)
	}()
}

func (g *Group) WhenDualComplete(completeFn func(first, second interface{})) {
	completeFn(<-g.results, <-g.results)
}

func (g *Group) WaitDual() (fist, second interface{}) {
	return <-g.results, <-g.results
}

func (g *Group) WhenDualCompleteAsync(completeFn func(first, second interface{})) {
	go func() {
		defer func() {
			if err := recover(); err != nil {
				loggers.Errorf("WhenDualCompleteAsync execute panic: %v, stack: %s", err, string(trace.Stack(2)))
			}
		}()
		g.WhenDualComplete(completeFn)
	}()
}

func (g *Group) WhenAllComplete(completeFn func(results ...interface{})) {
	var rs = make([]interface{}, g.size)
	for i := 0; i < g.size; i++ {
		rs[i] = <-g.results
	}
	completeFn(rs...)
}

func (g *Group) WhenAllCompleteAsync(completeFn func(results ...interface{})) {
	go func() {
		defer func() {
			if err := recover(); err != nil {
				loggers.Errorf("WhenAllCompleteAsync execute panic: %v, stack: %s", err, trace.Stack(2))
			}
		}()
		g.WhenAllComplete(completeFn)
	}()
}
